package com.xp.pgoj;

import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategy;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.atomic.AtomicBoolean;
import io.reactivex.Flowable;
/**
 * @author <a href=""https://github.com/Fzs0>逢左使</a>
 * @create 2024/10/01 00:03
 * @Description:
 */
public class GLMAITest {

    private static final String url = "https://open.bigmodel.cn/api/paas/v4/async-result/";

    // 静态 Map 变量
    private static final Map<String, Object> HEADERS = new HashMap<>();

    static {
        // 初始化静态 Map
        HEADERS.put("Content-Type", "application/json");
        HEADERS.put("Authorization", "Bearer 0315d4dab22f84e541d7046736c31dba.TsngMYRRhPCiLKu2");
    }
    private static final ObjectMapper mapper = defaultObjectMapper();
    public static ObjectMapper defaultObjectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
        mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE);
//        mapper.addMixIn(ChatFunction.class, ChatFunctionMixIn.class);
//        mapper.addMixIn(ChatCompletionRequest.class, ChatCompletionRequestMixIn.class);
//        mapper.addMixIn(ChatFunctionCall.class, ChatFunctionCallMixIn.class);
        return mapper;
    }

    public static void main(String[] args) {
        testSseInvoke();
    }
    
    /**
     * 同步调用
     */
    private static void testAsyncInvoke() {
        ClientV4 client = new ClientV4.Builder("0315d4dab22f84e541d7046736c31dba.TsngMYRRhPCiLKu2").build();
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "请你设计一道dfs的算法题");
        messages.add(chatMessage);

        long l = System.currentTimeMillis();
        String requestId = String.valueOf(l);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethodAsync)
                .requestId(requestId)
                .messages(messages)
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        System.out.println("modelAPIResponse output ID  ====>" + invokeModelApiResp.getData().getId());
        System.out.println("modelAPIResponse output  ====>" + invokeModelApiResp.getData());
        // 根据异步url来获取返回值
        String getResponseByAsyncUrl = url + String.valueOf(invokeModelApiResp.getData().getId());
        Timer timer = new Timer();
        timer.schedule(new TimerTask() {
            @Override
            public void run() {
                String response = HttpUtil.get(getResponseByAsyncUrl, HEADERS);
                handleResponse(response);
            }
        }, 0, 5000); // 每5秒检查一次
    }

    /**
     * 流式调用
     */
    private static void testSseInvoke() {
        ClientV4 client = new ClientV4.Builder("0315d4dab22f84e541d7046736c31dba.TsngMYRRhPCiLKu2").build();
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "请你设计一道dfs的算法题");
        messages.add(chatMessage);

        long l = System.currentTimeMillis();
        String requestId = String.valueOf(l);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethodAsync)
                .requestId(requestId)
                .messages(messages)
                .build();
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        if (sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            if (isFirst.getAndSet(false)) {
                                System.out.print("Response: ");
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                System.out.println("tool_calls: " + jsonString);
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                System.out.print(accumulator.getDelta().getContent());
                            }
                        }
                    })
                    .doOnComplete(System.out::println)
                    .lastElement()
                    .blockingGet();

            Choice choice = new Choice();
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
    }

    public static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> {
            return new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null, chunk.getChoices().get(0), chunk.getUsage(), chunk.getCreated(), chunk.getId());
        });
    }


    private static void handleResponse(String jsonResponse) {
        JSONObject jsonObject = new JSONObject(jsonResponse);
        String taskStatus = jsonObject.getStr("task_status");

        if ("SUCCESS".equals(taskStatus)) {
            String content = jsonObject.getJSONArray("choices").getJSONObject(0).getJSONObject("message").getStr("content");
            System.out.println("任务成功，内容是: " + content);

            // 停止监控
            System.exit(0);
        } else if ("FAILED".equals(taskStatus)) {
            System.out.println("任务失败，停止监控。");
            System.exit(0);
        } else {
            System.out.println("任务状态仍在处理中: " + taskStatus);
        }
    }
}
